# encoding: utf-8

import socket
from threading import Thread, Event

from util.log import getLogger
logger = getLogger(10)

from common import *  # BROKER_PORT
from protocal import PACK_HEAD_ERROR, PACK_HEAD_HEART

class UdpServer(Thread):
    def __init__(self):
        super().__init__()

        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.bind(("", BROKER_PORT))
        self.sock.settimeout(TIMEOUT_RECV)

        self.dict_clnt = {}  # address: {heart}

    def destroy(self):
        print("等待回收心跳线程")
        if hasattr(self, "tid_heart"):
            self.tid_heart.join()
        print("等待关闭socket")
        self.sock.close()

    def stop(self):
        self.isRunning = False

    def set_protocal(self, protocal):
        """ 载入数据处理器 """
        self.protocal = protocal

        def send_msg(msg, addr=None):
            if addr:
                self.sock.sendto(msg, addr)
            else:
                self.sock.send(msg)

        self.protocal.set_func_send(send_msg)

    def heart_echo(self):
        """ 启动线程实现对每个链接的心跳超时判断 """
        from time import sleep

        def timeout_count():
            while self.isRunning:
                sleep(HEART_LISTEN_PERIOD)

                list_closing = []
                for address, dict_clnt_info in self.dict_clnt.items():
                    dict_clnt_info["heart"] += HEART_LISTEN_PERIOD
                    if dict_clnt_info["heart"] > TIMEOUT_SHUTDOWN:
                        list_closing.append(address)

                for address in list_closing:
                    self.close_client(address)

        self.tid_heart = Thread(target=timeout_count)
        self.tid_heart.setDaemon(True)
        self.tid_heart.start()

    def listen(self):
        self.isRunning = True
        self.heart_echo()

        while self.isRunning:
            try:
                self.recv_msg()

            except socket.timeout:
                pass

            except ConnectionError:
                print("EEEExcept: 断开连接....")
                pass

            except Exception as e:
                # from traceback import print_exc
                # print_exc()
                logger.error(e)

    def run(self):
        self.listen()
        # 当监听结束时，也不再允许发送数据
        self.destroy()

    def close_client(self, address):
        logger.debug("关闭一个客户端连接")
        del self.dict_clnt[address]
        self.protocal.on_connect_close(address)

    def recv_msg(self):
        # 接收一整条数据
        data, addr = self.sock.recvfrom(BUFFER_SIZE)  # 超时
        # logger.debug(f"Recv Msg... [{data}] from [{addr}]")

        if addr not in self.dict_clnt:
            self.dict_clnt[addr] = {"heart": 0}
            logger.debug("新增客户端连接：【{}】".format(self.dict_clnt))

        # 区分指令、心跳、数据
        if data[:4] == PACK_HEAD_HEART:
            # logger.debug("接收到心跳")
            self.dict_clnt[addr]["heart"] = 0
            # heart_echo
            self.sock.sendto(PACK_HEAD_HEART, addr)
        elif data[:4] == PACK_HEAD_ERROR:
            reason = data[4:].decode()
            logger.error("Except: 【{}】".format(reason))
        else:
            self.protocal.reply(data, addr)


if __name__ == "__main__":
    from time import sleep
    from protocal import DpmDataTrans

    broker = UdpServer()
    protocal = DpmDataTrans()
    broker.set_protocal(protocal)
    broker.start()

    try:
        while True:
            data = input("Input the Data ([Q] for quit): ")
            if data == "Q":
                broker.stop()  # 关闭监听
                break
    except OSError:
        logger.info("由于心跳超时，已自动断开连接")

    broker.join()
